#include "ctracker.h"

namespace buola { namespace tracker {

static const float sScale=100;
static const int sNormalizeIter=5; //2 in original
static const float sEpsilon=0.001;
static const float sStdDev=0.02;

CTracker::CTracker(int K)
{
    mPriorDist=mat::CVec3f(sStdDev*sScale,sStdDev*sScale,sStdDev*sScale);
    mStdDev=mat::constant(K,3,sStdDev*sScale);
    mOutlierDist=mat::CVec3f(sStdDev*sScale,sStdDev*sScale,sStdDev*sScale);
    mOutlierStdDev=mat::CVec3f(sStdDev*sScale,sStdDev*sScale,sStdDev*sScale);
}

void CTracker::SetFeatures(mat::CMat_f &&pObject,mat::CMat_f &&pObserved)
{
    mObject=std::move(pObject);
    mObserved=std::move(pObserved);
    mVis=mat::ones(pObject.Rows());
}

void CTracker::ExpectationStep(const CTrackedObject &pObject)
{
    int K=mObject.Rows(); //number of nodes
    int N=mObserved.Rows(); //number of points
    int F=mObserved.Cols(); //number of features
    
    msg_info() << "exp step " << K << " " << N << " " << F << "\n";

    start_timer();
    mat::CMat_f lInvStdDev=1.0/mStdDev;
    mat::CMat_f lInvVar=sq(lInvStdDev);
    mat::CMat_f lSqDistsInvVar(K+1,N);
    
    for(int k=0;k<K;k++)
    {
        mat::CRow_f lTmp1=lInvVar(k,nAll);
        mat::CMat_f lTmp2=sq(mObserved-extend(mObject(k,nAll))).T(); //fix so it's easier
        mat::CRow_f lTmp3=lTmp1*lTmp2;
        lSqDistsInvVar(k,nAll)=lTmp3;
    }
    
    mat::CVec_f lOutlierInvStdDev=1.0/mOutlierStdDev;
    mat::CVec_f lOutlierInvVar=sq(lOutlierInvStdDev);
    lSqDistsInvVar(K,nAll)=dot(lOutlierInvVar,sq(mOutlierDist))*mat::CRow_f(mat::ones(1,N)); //fix mat so that is easier

    mat::CMat_f lPZGivenCExpPart=exp(-0.5*lSqDistsInvVar);
    mat::CVec_f lSqrtDetInvVar=prod(rows(lInvStdDev));

    mPZGivenC.Resize(K+1,N);
    for(int k=0;k<K;k++)  
    {
        mPZGivenC(k,nAll)=lSqrtDetInvVar[k]*lPZGivenCExpPart(k,nAll)*mVis[k];
    }
    mPZGivenC(K,nAll)=prod(lOutlierInvStdDev)*lPZGivenCExpPart(K,nAll);

    mPZGivenC=mPZGivenC**extend(1.0/sum(cols(mPZGivenC))); //check correct and check performance!!!

    end_timer();
    start_timer();

    //iteratively normalize rows and columns. columns are normalized to one and rows are normalized to the visibility term.
    for (int i=0;i<sNormalizeIter;i++) 
    {
        //normalize rows
        //FIXME does the k+1 row need to be normalized?
        for(int k=0;k<K;k++) 
        {
            if(mVis[k]<sEpsilon) 
                mPZGivenC(k,nAll)=mat::zeros(1,F);
            else
            {
                float lSum=sum(mPZGivenC(k,nAll));
                if(lSum!=0.0f)
                    mPZGivenC(k,nAll)=mPZGivenC(k,nAll)/lSum; //fix mat so that it's easier
            }
        }
        mPZGivenC(K,nAll)=mPZGivenC(K,nAll)/sum(mPZGivenC(K,nAll));
        //pZgivenC.row(K) /= pZgivenC.row(K).sum();

        //normalize cols
        mPZGivenC=mPZGivenC**extend(1.0/sum(cols(mPZGivenC)));
    }
    
    end_timer("E step");

    if(!all(is_finite(mPZGivenC)))
        throw XInvalid("not finite probability matrix");
}

void CTracker::MaximizationStep(CTrackedObject &pObject)
{
    start_timer();
    pObject.ApplyEvidence(mPZGivenC,mObserved,mObject);
    end_timer("M step");
}

/*namespace tracker*/ } /*namespace buola*/ }
